#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <errno.h>
#include "client.h"
#include "protocol.h"
#include "crypto/sm2_keyex.h"
#include "crypto/sm4_crypto.h"
#include "auth/sm3_auth.h"
#include "utils.h" // For get_file_size, etc.

#define BUFFER_SIZE 4096

// 初始化客户端
int client_init(void)
{
    // 初始化GmSSL
    if (init_gmssl() != 0)
    {
        fprintf(stderr, "初始化GmSSL引擎失败\n");
        return -1;
    }

    return 0;
}

// 初始化会话
int client_session_init(client_session_t *session)
{
    if (!session)
    {
        return -1;
    }

    // 初始化会话信息
    memset(session, 0, sizeof(client_session_t));
    session->sock = -1;
    session->state = SESSION_STATE_DISCONNECTED;
    session->current_seq = 1; // 从1开始递增

    // 初始化当前传输状态
    session->current_transfer.state = TRANSFER_STATE_NONE;
    session->current_transfer.direction = TRANSFER_DIRECTION_NONE;

    // 设置一个合理的超时重试值
    session->reconnect_attempts = 0;
    session->max_reconnect_attempts = 3;

    return 0;
}

// 连接服务器
int client_connect(client_session_t *session, const char *server_ip, int port)
{
    if (!session || !server_ip || port <= 0)
    {
        return -1;
    }

    // 如果已连接，先断开
    if (session->sock >= 0)
    {
        close(session->sock);
        session->sock = -1;
    }

    // 重置会话状态
    session->state = SESSION_STATE_DISCONNECTED;

    // 创建套接字
    session->sock = socket(AF_INET, SOCK_STREAM, 0);
    if (session->sock < 0)
    {
        perror("创建套接字失败");
        return -1;
    }

    // 设置连接超时
    struct timeval timeout;
    timeout.tv_sec = 10; // 10秒超时
    timeout.tv_usec = 0;
    if (setsockopt(session->sock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0)
    {
        perror("设置发送超时失败");
        close(session->sock);
        session->sock = -1;
        return -1;
    }

    if (setsockopt(session->sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0)
    {
        perror("设置接收超时失败");
        close(session->sock);
        session->sock = -1;
        return -1;
    }

    // 连接服务器
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(server_ip);
    server_addr.sin_port = htons(port);

    if (connect(session->sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0)
    {
        perror("连接服务器失败");
        close(session->sock);
        session->sock = -1;
        return -1;
    }

    // 更新会话状态
    session->state = SESSION_STATE_CONNECTED;

    // 生成本地SM2密钥对
    if (sm2_generate_key_pair(&session->local_key) != 0)
    {
        fprintf(stderr, "生成SM2密钥对失败\n");
        close(session->sock);
        session->sock = -1;
        session->state = SESSION_STATE_DISCONNECTED;
        return -1;
    }

    return 0;
}

// 执行SM2密钥交换
int client_key_exchange(client_session_t *session)
{
    if (!session || session->sock < 0 || session->state != SESSION_STATE_CONNECTED)
    {
        return -1;
    }

    // 发送本地公钥
    key_exchange_t local_key;
    if (sm2_get_public_key_coords(&session->local_key,
                                  local_key.public_key_x, sizeof(local_key.public_key_x),
                                  local_key.public_key_y, sizeof(local_key.public_key_y)) != 0)
    {
        fprintf(stderr, "获取本地公钥坐标失败\n");
        client_disconnect(session);
        return -1;
    }

    uint16_t seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_KEY_EXCHANGE, seq, &local_key, sizeof(local_key)) != 0)
    {
        fprintf(stderr, "发送密钥交换消息失败\n");
        client_disconnect(session);
        return -1;
    }

    // 接收服务器公钥
    uint8_t type;
    uint16_t recv_seq;
    void *data = NULL;
    uint32_t len = 0;

    int recv_ret_keyex = recv_specific_message(session->sock, MSG_TYPE_KEY_EXCHANGE, 0, &data, &len);

    if (recv_ret_keyex != 0 || data == NULL)
    {
        fprintf(stderr, "接收密钥交换响应失败\n");
        client_disconnect(session);
        return -1;
    }

    if (len != sizeof(key_exchange_t))
    {
        fprintf(stderr, "接收到的密钥交换数据长度不正确\n");
        free(data);
        client_disconnect(session);
        return -1;
    }

    // 保存服务器公钥
    key_exchange_t *remote_key = (key_exchange_t *)data;
    memcpy(session->remote_public_x, remote_key->public_key_x, sizeof(session->remote_public_x));
    memcpy(session->remote_public_y, remote_key->public_key_y, sizeof(session->remote_public_y));

    // 派生共享密钥
    if (sm2_derive_shared_key(&session->local_key,
                              session->remote_public_x, sizeof(session->remote_public_x),
                              session->remote_public_y, sizeof(session->remote_public_y),
                              session->shared_key, sizeof(session->shared_key)) != 0)
    {
        fprintf(stderr, "派生共享密钥失败\n");
        free(data);
        client_disconnect(session);
        return -1;
    }

    free(data);
    session->state = SESSION_STATE_KEY_EXCHANGED;
    return 0;
}

// 用户认证
int client_authenticate(client_session_t *session, const char *username, const char *password)
{
    if (!session || !username || !password || session->sock < 0 || session->state != SESSION_STATE_KEY_EXCHANGED)
    {
        fprintf(stderr, "[Client Authenticate] Error: Invalid session state (%d) or parameters. Expected KEY_EXCHANGED (%d).\n", session ? session->state : -1, SESSION_STATE_KEY_EXCHANGED);
        return -1;
    }

    // 准备认证请求
    auth_request_t auth_req;
    memset(&auth_req, 0, sizeof(auth_req));

    // 拷贝用户名（防止缓冲区溢出）
    strncpy(auth_req.username, username, sizeof(auth_req.username) - 1);
    auth_req.username[sizeof(auth_req.username) - 1] = '\0';

    // 拷贝密码（实际应用中应该使用哈希值，但这里简化处理）
    strncpy(auth_req.password, password, sizeof(auth_req.password) - 1);
    auth_req.password[sizeof(auth_req.password) - 1] = '\0';

    // 发送认证请求
    uint16_t seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_AUTH_REQUEST, seq, &auth_req, sizeof(auth_req)) != 0)
    {
        fprintf(stderr, "发送认证请求失败\n");
        client_disconnect(session);
        return -1;
    }

    // 接收认证响应
    void *data = NULL;
    uint32_t len = 0;
    int recv_ret_auth = recv_specific_message(session->sock, MSG_TYPE_AUTH_RESPONSE, 0, &data, &len);

    if (recv_ret_auth != 0 || data == NULL)
    {
        fprintf(stderr, "接收认证响应失败\n");
        client_disconnect(session);
        return -1;
    }

    if (len != sizeof(auth_response_t))
    {
        fprintf(stderr, "接收到的认证响应数据长度不正确\n");
        free(data);
        client_disconnect(session);
        return -1;
    }

    // 解析认证响应
    auth_response_t *auth_resp = (auth_response_t *)data;
    if (auth_resp->result != ERR_SUCCESS)
    {
        fprintf(stderr, "认证失败，错误码: %d\n", auth_resp->result);
        free(data);
        return -1;
    }

    // 保存认证信息
    session->user_id = auth_resp->user_id;
    strncpy(session->username, username, sizeof(session->username) - 1);
    session->username[sizeof(session->username) - 1] = '\0';
    memcpy(session->token, auth_resp->token, sizeof(session->token));

    // 更新会话状态
    session->state = SESSION_STATE_AUTHENTICATED;

    free(data);
    return 0;
}

// 获取文件列表
int client_get_file_list(client_session_t *session, const char *path, void **file_list, int *file_count)
{
    if (!session || !path || !file_list || !file_count ||
        session->sock < 0 || session->state != SESSION_STATE_AUTHENTICATED)
    {
        return -1;
    }

    // 初始化输出参数
    *file_list = NULL;
    *file_count = 0;

    // 准备文件列表请求
    file_list_request_t list_req;
    memset(&list_req, 0, sizeof(list_req));

    // 设置文件路径
    strncpy(list_req.path, path, sizeof(list_req.path) - 1);
    list_req.path[sizeof(list_req.path) - 1] = '\0';

    // 设置会话令牌
    memcpy(list_req.token, session->token, sizeof(list_req.token));

    // 发送文件列表请求
    uint16_t seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_FILE_LIST_REQUEST, seq, &list_req, sizeof(list_req)) != 0)
    {
        fprintf(stderr, "发送文件列表请求失败\n");
        return -1;
    }

    // 接收文件列表响应
    void *data = NULL;
    uint32_t len = 0;
    int recv_ret_list = recv_specific_message(session->sock, MSG_TYPE_FILE_LIST_RESPONSE, 0, &data, &len);

    if (recv_ret_list != 0 || data == NULL)
    {
        fprintf(stderr, "接收文件列表响应失败\n");
        return -1;
    }

    // 解析文件列表
    if (len < sizeof(uint32_t))
    {
        fprintf(stderr, "接收到的文件列表数据长度不正确\n");
        free(data);
        return -1;
    }

    // 获取文件数量
    uint32_t count = *((uint32_t *)data);

    // 确保数据长度合理
    if (len != sizeof(uint32_t) + count * sizeof(file_entry_t))
    {
        fprintf(stderr, "文件列表数据长度与文件数量不符\n");
        free(data);
        return -1;
    }

    // 分配内存保存文件条目
    file_entry_t *entries = malloc(count * sizeof(file_entry_t));
    if (!entries)
    {
        fprintf(stderr, "内存分配失败\n");
        free(data);
        return -1;
    }

    // 复制文件条目
    file_entry_t *src_entries = (file_entry_t *)((char *)data + sizeof(uint32_t));
    memcpy(entries, src_entries, count * sizeof(file_entry_t));

    // 设置返回值
    *file_list = entries;
    *file_count = count;

    free(data);
    return 0;
}

// 上传文件
int client_upload_file(client_session_t *session, const char *local_path, const char *remote_path)
{
    if (!session || session->state != SESSION_STATE_AUTHENTICATED)
    {
        fprintf(stderr, "错误: 请先成功登录 (当前会话状态: %d). 上传需要认证会话.\n", session ? session->state : -1);
        return -1;
    }

    FILE *file = fopen(local_path, "rb");
    if (!file)
    {
        perror("打开本地文件失败");
        return -1;
    }

    uint64_t file_size = get_file_size(local_path);
    if (file_size == (uint64_t)-1)
    {
        fprintf(stderr, "获取文件大小失败: %s\n", local_path);
        fclose(file);
        return -1;
    }

    // 准备文件上传
    printf("正在上传文件 %s -> %s (%llu bytes)...\n", local_path, remote_path, (unsigned long long)file_size);

    // 初始化传输状态
    session->current_transfer.state = TRANSFER_STATE_STARTING;
    session->current_transfer.direction = TRANSFER_DIRECTION_UPLOAD;
    session->current_transfer.total_size = file_size;
    session->current_transfer.transferred_size = 0;
    strncpy(session->current_transfer.src_path, local_path, sizeof(session->current_transfer.src_path) - 1);
    strncpy(session->current_transfer.dst_path, remote_path, sizeof(session->current_transfer.dst_path) - 1);

    // 1. 发送文件上传请求
    file_upload_request_t upload_req;
    strncpy(upload_req.path, remote_path, sizeof(upload_req.path) - 1);
    upload_req.path[sizeof(upload_req.path) - 1] = '\0';
    upload_req.size = file_size;
    memcpy(upload_req.token, session->token, sizeof(session->token));

    uint16_t upload_req_seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_FILE_UPLOAD_REQUEST, upload_req_seq, &upload_req, sizeof(upload_req)) != 0)
    {
        fprintf(stderr, "发送文件上传请求失败\n");
        fclose(file);
        return -1;
    }

    // 2. 接收文件上传响应
    uint8_t resp_type;
    uint16_t resp_seq;
    void *resp_data = NULL;
    uint32_t resp_len;

    int recv_ret_upload_init = recv_message(session->sock, &resp_type, &resp_seq, &resp_data, &resp_len);

    if (recv_ret_upload_init != 0)
    {
        fprintf(stderr, "接收上传响应失败\n");
        fclose(file);
        if (resp_data)
            free(resp_data);
        return -1;
    }

    if (resp_type == MSG_TYPE_ERROR && resp_data)
    {
        error_message_t *err_msg = (error_message_t *)resp_data;
        fprintf(stderr, "错误: %s (代码: %d)\n", err_msg->message, err_msg->code);
        fclose(file);
        free(resp_data);
        return -1;
    }

    if (resp_type != MSG_TYPE_FILE_UPLOAD_RESPONSE || !resp_data || resp_len < sizeof(file_upload_response_t))
    {
        fprintf(stderr, "接收到无效的文件上传响应 (type: %02X, len: %u)\n", resp_type, resp_len);
        fclose(file);
        if (resp_data)
            free(resp_data);
        return -1;
    }

    file_upload_response_t *upload_resp = (file_upload_response_t *)resp_data;
    if (upload_resp->result != ERR_SUCCESS)
    {
        fprintf(stderr, "文件上传被服务器拒绝 (代码: %d)\n", upload_resp->result);
        fclose(file);
        free(resp_data);
        return -1;
    }

    printf("服务器接受上传请求. Upload ID: %u, Suggested block size: %u\n", upload_resp->upload_id, upload_resp->block_size);
    uint32_t server_upload_id = upload_resp->upload_id;

    // 更新传输状态
    session->current_transfer.transfer_id = server_upload_id;
    session->current_transfer.state = TRANSFER_STATE_ACTIVE;
    session->current_transfer.block_size = upload_resp->block_size;

    free(resp_data);
    resp_data = NULL;

    const uint32_t CLIENT_BLOCK_PAYLOAD_SIZE = 4080; // Max plaintext bytes per data message block
    unsigned char read_buffer[CLIENT_BLOCK_PAYLOAD_SIZE];
    // Max encrypted size: CLIENT_BLOCK_PAYLOAD_SIZE + SM4_BLOCK_SIZE (for padding)
    // Max data in message: SM4_IV_SIZE (for first block) + encrypted size
    unsigned char sm4_payload_buffer[SM4_IV_SIZE + CLIENT_BLOCK_PAYLOAD_SIZE + SM4_BLOCK_SIZE];
    unsigned char iv[SM4_IV_SIZE];

    // Generate IV for this upload session
    if (generate_sm4_iv(iv) != 0)
    { // Assumes generate_sm4_iv is from your sm4_crypto.h
        fprintf(stderr, "错误: 生成SM4 IV失败\n");
        fclose(file);
        return -1;
    }

    // 3. 发送文件数据块
    uint32_t block_seq_counter = 0;
    size_t bytes_read_from_file;
    uint64_t total_bytes_sent_for_file = 0;
    int is_first_data_block = 1;

    while (1)
    {
        bytes_read_from_file = fread(read_buffer, 1, CLIENT_BLOCK_PAYLOAD_SIZE, file);

        if (bytes_read_from_file == 0 && feof(file))
        {
            // End of file, all data has been processed in previous iterations if any.
            // If file_size is 0, this loop won't execute, which is fine.
            // If last read was exactly CLIENT_BLOCK_PAYLOAD_SIZE, we still need to send one last (empty or padded) encrypted block.
            // The sm4_cbc_encrypt_data should handle zero-length input for final padding if needed.
            // However, our current loop structure sends data only if bytes_read_from_file > 0.
            // For simplicity, we will assume that if bytes_read_from_file is 0 here, we are done.
            // A more robust solution handles a final encryption call even if last read is 0 to ensure padding is sent.
            break;
        }
        if (bytes_read_from_file == 0 && ferror(file))
        {
            perror("读取本地文件时发生错误");
            fclose(file);
            return -1;
        }

        size_t encrypted_len = sizeof(sm4_payload_buffer) - (is_first_data_block ? SM4_IV_SIZE : 0);
        unsigned char *encryption_output_ptr = sm4_payload_buffer;
        if (is_first_data_block)
        {
            memcpy(encryption_output_ptr, iv, SM4_IV_SIZE); // Prepend IV
            encryption_output_ptr += SM4_IV_SIZE;
        }

        // Encrypt the actual data block
        if (sm4_cbc_encrypt_data(read_buffer, bytes_read_from_file,
                                 session->shared_key, iv,
                                 encryption_output_ptr, &encrypted_len) != 0)
        {
            fprintf(stderr, "错误: 文件块加密失败 (block %u)\n", block_seq_counter);
            fclose(file);
            return -1;
        }

        uint32_t data_payload_len_to_send = encrypted_len;
        if (is_first_data_block)
        {
            data_payload_len_to_send += SM4_IV_SIZE; // Total length includes IV for the first block
        }

        size_t data_msg_header_size = offsetof(file_upload_data_t, data);
        size_t total_data_msg_len = data_msg_header_size + data_payload_len_to_send;
        file_upload_data_t *data_msg_to_send = (file_upload_data_t *)malloc(total_data_msg_len);
        if (!data_msg_to_send)
        {
            perror("Failed to allocate memory for upload data message");
            fclose(file);
            return -1;
        }

        data_msg_to_send->upload_id = server_upload_id;
        data_msg_to_send->block_seq = block_seq_counter;
        data_msg_to_send->data_len = data_payload_len_to_send; // This is IV + Ciphertext or just Ciphertext
        memcpy(data_msg_to_send->data, sm4_payload_buffer, data_payload_len_to_send);

        uint16_t upload_data_seq = session->current_seq++;
        if (send_message(session->sock, MSG_TYPE_FILE_UPLOAD_DATA, upload_data_seq, data_msg_to_send, total_data_msg_len) != 0)
        {
            fprintf(stderr, "发送文件数据块 %u 失败\n", block_seq_counter);
            free(data_msg_to_send);
            fclose(file);
            return -1;
        }
        free(data_msg_to_send);
        is_first_data_block = 0; // IV is only sent with the first block

        // 4. 等待服务器对该数据块的响应 (MSG_TYPE_FILE_UPLOAD_DATA_RESPONSE)
        int recv_ret_upload_ack = recv_message(session->sock, &resp_type, &resp_seq, &resp_data, &resp_len);
        if (recv_ret_upload_ack != 0)
        {
            fprintf(stderr, "接收数据块 %u 响应失败\n", block_seq_counter);
            fclose(file);
            if (resp_data)
                free(resp_data);
            return -1;
        }

        if (resp_type == MSG_TYPE_ERROR && resp_data)
        {
            error_message_t *err_msg = (error_message_t *)resp_data;
            fprintf(stderr, "服务器错误 (块 %u): %s (代码: %d)\n", block_seq_counter, err_msg->message, err_msg->code);
            fclose(file);
            free(resp_data);
            return -1;
        }

        if (resp_type != MSG_TYPE_FILE_UPLOAD_DATA_RESPONSE || !resp_data || resp_len < sizeof(file_upload_data_response_t))
        {
            fprintf(stderr, "接收到无效的数据块响应 (块 %u, type: %02X, len: %u)\n", block_seq_counter, resp_type, resp_len);
            fclose(file);
            if (resp_data)
                free(resp_data);
            return -1;
        }

        file_upload_data_response_t *ack_from_server = (file_upload_data_response_t *)resp_data;
        if (ack_from_server->result != ERR_SUCCESS)
        {
            fprintf(stderr, "服务器未能成功接收数据块 %u (代码: %d)\n", ack_from_server->block_seq, ack_from_server->result);
            fclose(file);
            free(resp_data);
            return -1;
        }

        if (ack_from_server->upload_id != server_upload_id || ack_from_server->block_seq != block_seq_counter)
        {
            fprintf(stderr, "数据块响应的 upload_id 或 block_seq 不匹配 (期望 %u/%u,收到 %u/%u )\n",
                    server_upload_id, block_seq_counter, ack_from_server->upload_id, ack_from_server->block_seq);
            fclose(file);
            free(resp_data);
            return -1;
        }

        // total_bytes_sent_for_file should track original file bytes, not encrypted bytes
        total_bytes_sent_for_file += bytes_read_from_file;

        // 更新传输进度
        session->current_transfer.transferred_size = total_bytes_sent_for_file;

        printf("数据块 %u (明文 %zu B) 发送并确认成功. 已发送: %llu / %llu bytes\n",
               block_seq_counter, bytes_read_from_file, (unsigned long long)total_bytes_sent_for_file, (unsigned long long)file_size);
        block_seq_counter++;
        free(resp_data);
        resp_data = NULL;

        if (total_bytes_sent_for_file >= file_size)
        {
            // All actual file bytes have been read and processed for encryption.
            // The last call to sm4_cbc_encrypt_data would have included padding if necessary.
            break;
        }
    }

    if (ferror(file))
    {
        perror("读取本地文件时发生错误");
        fclose(file);
        return -1;
    }

    fclose(file);
    if (total_bytes_sent_for_file == file_size)
    {
        printf("文件 %s 上传完成 (%llu bytes).\n", local_path, (unsigned long long)total_bytes_sent_for_file);
        session->current_transfer.state = TRANSFER_STATE_COMPLETED;
        return 0;
    }
    else
    {
        fprintf(stderr, "文件 %s 上传可能未完成. 发送字节 %llu 与文件大小 %llu 不符.\n",
                local_path, (unsigned long long)total_bytes_sent_for_file, (unsigned long long)file_size);
        session->current_transfer.state = TRANSFER_STATE_ERROR;
        session->current_transfer.error_code = ERR_INTERNAL_ERROR;
        return -1;
    }
}

// 下载文件
int client_download_file(client_session_t *session, const char *remote_path, const char *local_path)
{
    if (!session || session->sock < 0 ||
        session->state != SESSION_STATE_AUTHENTICATED ||
        !remote_path || !local_path)
    {
        return -1;
    }

    // 检查是否有其他传输正在进行
    if (session->current_transfer.state == TRANSFER_STATE_ACTIVE)
    {
        fprintf(stderr, "已有传输任务正在进行\n");
        return -1;
    }

    // 准备下载请求
    file_download_request_t download_req;
    memset(&download_req, 0, sizeof(download_req));
    strncpy(download_req.path, remote_path, sizeof(download_req.path) - 1);
    memcpy(download_req.token, session->token, sizeof(download_req.token));

    fprintf(stdout, "客户端：向服务端发送文件 '%s' 的下载请求...\\n", remote_path);
    // 发送下载请求
    uint16_t seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_FILE_DOWNLOAD_REQUEST, seq, &download_req, sizeof(download_req)) != 0)
    {
        fprintf(stderr, "发送文件下载请求失败\n");
        fprintf(stdout, "客户端：向服务端发送文件 '%s' 下载请求失败。\n", remote_path);
        return -1;
    }

    // 接收下载响应
    void *data = NULL;
    uint32_t len = 0;

    int recv_ret_dl_init = recv_specific_message(session->sock, MSG_TYPE_FILE_DOWNLOAD_RESPONSE, 0, &data, &len);

    if (recv_ret_dl_init != 0 || data == NULL)
    {
        fprintf(stderr, "接收下载响应失败\n");
        fprintf(stdout, "客户端：接收服务端对文件 '%s' 下载请求的响应失败。\n", remote_path);
        return -1;
    }

    if (len != sizeof(file_download_response_t))
    {
        fprintf(stderr, "接收到的下载响应数据长度不正确\n");
        free(data);
        return -1;
    }

    // 处理下载响应
    file_download_response_t *download_resp = (file_download_response_t *)data;
    if (download_resp->result != ERR_SUCCESS)
    {
        fprintf(stderr, "下载请求被拒绝，错误码: %d\n", download_resp->result);
        fprintf(stdout, "客户端：文件 '%s' 下载请求被服务端拒绝，错误码: %d。\n", remote_path, download_resp->result);
        free(data);
        return -1;
    }
    fprintf(stdout, "客户端：服务端接受文件 '%s' 的下载请求。文件大小: %llu B, 下载ID: %u。\n", remote_path, (unsigned long long)download_resp->size, download_resp->download_id);

    // 创建本地文件
    int local_fd = open(local_path, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (local_fd < 0)
    {
        perror("创建本地文件失败");
        free(data);
        return -1;
    }

    // 保存传输信息
    session->current_transfer.transfer_id = download_resp->download_id;
    session->current_transfer.direction = TRANSFER_DIRECTION_DOWNLOAD;
    session->current_transfer.state = TRANSFER_STATE_ACTIVE;
    session->current_transfer.total_size = download_resp->size;
    session->current_transfer.transferred_size = 0;
    session->current_transfer.block_size = download_resp->block_size;
    session->current_transfer.current_block = 0;
    strncpy(session->current_transfer.src_path, remote_path, sizeof(session->current_transfer.src_path) - 1);
    strncpy(session->current_transfer.dst_path, local_path, sizeof(session->current_transfer.dst_path) - 1);

    free(data);

    // 密钥材料
    unsigned char key[SM4_KEY_SIZE];
    unsigned char iv[SM4_IV_SIZE];

    // 根据共享密钥派生解密密钥
    memcpy(key, session->shared_key, SM4_KEY_SIZE);

    // 下载文件数据
    int block_num = 0;
    uint64_t total_received = 0;

    while (total_received < session->current_transfer.total_size)
    {
        // 发送文件数据块请求
        file_download_data_request_t data_request;
        data_request.download_id = session->current_transfer.transfer_id;
        data_request.block_num = block_num;

        // fprintf(stderr, "[CLIENT_DOWNLOAD] 发送数据块请求: ID=%u, 块号=%u\n",
        //         data_request.download_id, data_request.block_num);
        fprintf(stdout, "客户端：向服务端请求文件 '%s' (下载ID: %u) 的数据块 %d...\\n", session->current_transfer.src_path, data_request.download_id, data_request.block_num);

        uint16_t req_seq = session->current_seq++;
        if (send_message(session->sock, MSG_TYPE_FILE_DOWNLOAD_DATA_REQUEST, req_seq, &data_request, sizeof(data_request)) != 0)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 发送数据块请求失败\n");
            fprintf(stdout, "客户端：向服务端请求文件 '%s' (下载ID: %u) 的数据块 %d 失败。\n", session->current_transfer.src_path, data_request.download_id, data_request.block_num);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_CONNECTION_ERROR;
            return -1;
        }

        // 接收文件数据
        void *block_data = NULL;
        uint32_t block_len = 0;

        // 添加重试逻辑
        int retry_count = 0;
        const int max_retries = 3;
        int recv_ret = -1;

        while (retry_count < max_retries)
        {
            recv_ret = recv_specific_message(session->sock, MSG_TYPE_FILE_DOWNLOAD_DATA, req_seq, &block_data, &block_len);
            if (recv_ret == 0)
            {
                // 成功接收到响应
                break;
            }
            else if (recv_ret == ERR_RECV_WOULD_BLOCK)
            {
                // 非阻塞socket暂时没有数据，等待后重试
                fprintf(stderr, "[CLIENT_DOWNLOAD] 接收响应时遇到WOULD_BLOCK，重试 (%d/%d)...\n",
                        retry_count + 1, max_retries);

                // 使用简单循环延时替代usleep
                for (volatile int i = 0; i < 10000000; i++)
                {
                }

                retry_count++;
                continue;
            }
            else
            {
                // 其他错误，直接跳出
                break;
            }
        }

        if (recv_ret != 0 || block_data == NULL)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 接收文件数据失败，错误码: %d，重试次数: %d\n",
                    recv_ret, retry_count);
            fprintf(stdout, "客户端：从服务端接收文件 '%s' (下载ID: %u) 的数据块 %d 失败。\n",
                    session->current_transfer.src_path, session->current_transfer.transfer_id, block_num);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_CONNECTION_ERROR;
            return -1;
        }

        // 验证数据结构
        if (block_len < sizeof(file_download_data_t))
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 接收到的文件数据长度不正确: %u < %zu\n",
                    block_len, sizeof(file_download_data_t));
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_PROTOCOL_VERSION;
            return -1;
        }

        // 处理文件数据
        file_download_data_t *download_data = (file_download_data_t *)block_data;

        // 验证下载ID和块序号
        if (download_data->download_id != session->current_transfer.transfer_id)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 下载ID不匹配: 预期 %u, 收到 %u\n",
                    session->current_transfer.transfer_id, download_data->download_id);
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_PROTOCOL_VERSION;
            return -1;
        }

        if (download_data->block_num != block_num)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 块序号不匹配: 预期 %d, 收到 %d\n",
                    block_num, download_data->block_num);
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_PROTOCOL_VERSION;
            return -1;
        }

        fprintf(stdout, "客户端：成功从服务端接收到文件 '%s' (下载ID: %u) 的加密数据块 %d，密文长度: %u B。\n",
                session->current_transfer.src_path, download_data->download_id, block_num, download_data->data_len);

        // 解析数据部分
        unsigned char *encrypted_data = download_data->data;
        size_t encrypted_len = download_data->data_len;

        // 提取IV（位于加密数据的开头）
        if (encrypted_len < SM4_IV_SIZE)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 加密数据长度过短，无法提取IV: %zu < %d\n",
                    encrypted_len, SM4_IV_SIZE);
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_PROTOCOL_VERSION;
            return -1;
        }

        memcpy(iv, encrypted_data, SM4_IV_SIZE);
        fprintf(stdout, "客户端：从接收到的数据块 %d 中提取SM4-CBC模式的IV。\n", block_num);

        // 解密数据
        unsigned char *decrypted_data = (unsigned char *)malloc(encrypted_len);
        if (!decrypted_data)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 内存分配失败\n");
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_INTERNAL_ERROR;
            return -1;
        }

        size_t decrypted_len;
        fprintf(stdout, "客户端：准备使用SM4-CBC模式解密数据块 %d (密文部分长度: %zu B)。\n", block_num, encrypted_len - SM4_IV_SIZE);
        if (sm4_cbc_decrypt_data(encrypted_data + SM4_IV_SIZE, encrypted_len - SM4_IV_SIZE,
                                 key, iv, decrypted_data, &decrypted_len) != 0)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 解密数据失败\n");
            fprintf(stdout, "客户端：SM4解密数据块 %d (文件 '%s') 失败。\n", block_num, session->current_transfer.src_path);
            free(decrypted_data);
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_CRYPTO_ERROR;
            return -1;
        }

        fprintf(stdout, "客户端：SM4解密数据块 %d (文件 '%s') 成功，得到明文数据 %zu B。\n", block_num, session->current_transfer.src_path, decrypted_len);

        // 写入解密后的数据到文件
        ssize_t written = write(local_fd, decrypted_data, decrypted_len);
        if (written != decrypted_len)
        {
            perror("[CLIENT_DOWNLOAD] 写入文件失败");
            free(decrypted_data);
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_INTERNAL_ERROR;
            return -1;
        }

        fprintf(stdout, "客户端：已将解密后的数据块 %d (%zd B) 写入本地文件 '%s'。\n", block_num, written, local_path);
        free(decrypted_data);

        // 发送下载确认
        file_download_ack_t download_ack;
        download_ack.download_id = session->current_transfer.transfer_id;
        download_ack.block_num = block_num;
        download_ack.result = ERR_SUCCESS;

        uint16_t ack_seq = session->current_seq++;
        if (send_message(session->sock, MSG_TYPE_FILE_DOWNLOAD_ACK, ack_seq, &download_ack, sizeof(download_ack)) != 0)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 发送下载确认失败\n");
            free(block_data);
            close(local_fd);
            session->current_transfer.state = TRANSFER_STATE_ERROR;
            session->current_transfer.error_code = ERR_CONNECTION_ERROR;
            return -1;
        }

        // 更新传输状态
        total_received += decrypted_len;
        session->current_transfer.transferred_size = total_received;
        session->current_transfer.current_block = block_num;

        // 检查是否是最后一个数据块
        fprintf(stderr, "[CLIENT_DOWNLOAD] 已接收 %llu/%llu 字节，当前块: %d，是否为最后块: %d\n",
                (unsigned long long)total_received, (unsigned long long)session->current_transfer.total_size,
                block_num, download_data->is_last);

        free(block_data);
        block_data = NULL;

        if (download_data->is_last)
        {
            fprintf(stderr, "[CLIENT_DOWNLOAD] 收到最后一个数据块标记，结束下载\n");
            break;
        }

        // 继续下载下一块
        block_num++;
    }

    close(local_fd);

    if (total_received == session->current_transfer.total_size ||
        (session->current_transfer.total_size == 0 && total_received >= 0))
    {
        // 下载完成
        session->current_transfer.state = TRANSFER_STATE_COMPLETED;
        fprintf(stderr, "[CLIENT_DOWNLOAD] 文件下载完成: %s (%llu 字节)\n",
                local_path, (unsigned long long)total_received);
        fprintf(stdout, "客户端：文件 '%s' (大小: %llu B) 已成功下载并完成SM4解密。\n", local_path, (unsigned long long)total_received);
        return 0;
    }
    else
    {
        // 下载不完整
        session->current_transfer.state = TRANSFER_STATE_ERROR;
        session->current_transfer.error_code = ERR_PROTOCOL_VERSION;
        fprintf(stderr, "[CLIENT_DOWNLOAD] 下载不完整: 预期 %llu 字节, 实际接收 %llu 字节\n",
                (unsigned long long)session->current_transfer.total_size,
                (unsigned long long)total_received);
        fprintf(stdout, "客户端：文件 '%s' 下载失败或未完整接收。预期大小: %llu B, 实际接收明文大小: %llu B。\n",
                local_path, (unsigned long long)session->current_transfer.total_size, (unsigned long long)total_received);
        return -1;
    }
}

// 取消文件传输
int client_cancel_transfer(client_session_t *session)
{
    if (!session || session->sock < 0 ||
        session->state != SESSION_STATE_AUTHENTICATED)
    {
        return -1;
    }

    // 如果当前没有传输在进行中，返回成功
    if (session->current_transfer.state != TRANSFER_STATE_ACTIVE)
    {
        return 0; // 没有传输需要取消，视为成功
    }

    // 创建取消传输请求消息
    transfer_cancel_request_t cancel_req;
    memset(&cancel_req, 0, sizeof(cancel_req));

    // 填充取消请求数据
    cancel_req.transfer_id = session->current_transfer.transfer_id;
    if (session->current_transfer.direction == TRANSFER_DIRECTION_UPLOAD)
    {
        cancel_req.type = TRANSFER_TYPE_UPLOAD;
    }
    else if (session->current_transfer.direction == TRANSFER_DIRECTION_DOWNLOAD)
    {
        cancel_req.type = TRANSFER_TYPE_DOWNLOAD;
    }
    else
    {
        fprintf(stderr, "[CLIENT_CANCEL] 未知的传输方向: %d\n", session->current_transfer.direction);
        return -1;
    }

    // 发送取消请求
    uint16_t seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_TRANSFER_CANCEL_REQUEST, seq, &cancel_req, sizeof(cancel_req)) != 0)
    {
        return -1;
    }

    // 等待取消响应
    void *resp_data = NULL;
    uint32_t resp_len = 0;
    int recv_ret_cancel = recv_specific_message(session->sock, MSG_TYPE_TRANSFER_CANCEL_RESPONSE, seq, &resp_data, &resp_len);

    if (recv_ret_cancel != 0 || !resp_data)
    {
        // 即使未收到响应，也将传输标记为已取消
        session->current_transfer.state = TRANSFER_STATE_CANCELLED;
        return -1;
    }

    // 解析取消响应
    if (resp_len < sizeof(transfer_cancel_response_t))
    {
        fprintf(stderr, "[CLIENT_CANCEL] 接收到的取消响应数据长度不正确: %u\n", resp_len);
        free(resp_data);
        session->current_transfer.state = TRANSFER_STATE_CANCELLED;
        return -1;
    }

    transfer_cancel_response_t *cancel_resp = (transfer_cancel_response_t *)resp_data;
    int result = (cancel_resp->result == ERR_SUCCESS) ? 0 : -1;

    if (result == 0)
    {
        fprintf(stderr, "[CLIENT_CANCEL] 传输取消成功: ID=%u\n", session->current_transfer.transfer_id);
    }
    else
    {
        fprintf(stderr, "[CLIENT_CANCEL] 传输取消失败: ID=%u, 错误码=%d\n",
                session->current_transfer.transfer_id, cancel_resp->result);
    }

    free(resp_data);

    // 重置传输状态
    session->current_transfer.state = TRANSFER_STATE_CANCELLED;

    return result;
}

// 获取文件传输进度
int client_get_transfer_progress(client_session_t *session, uint64_t *transferred, uint64_t *total)
{
    if (!session || !transferred || !total)
    {
        return -1;
    }

    *transferred = session->current_transfer.transferred_size;
    *total = session->current_transfer.total_size;

    return session->current_transfer.state;
}

// 断开连接
int client_disconnect(client_session_t *session)
{
    if (!session)
    {
        return -1;
    }

    // 如果有活跃的传输任务，先取消
    if (session->current_transfer.state == TRANSFER_STATE_ACTIVE ||
        session->current_transfer.state == TRANSFER_STATE_STARTING)
    {
        client_cancel_transfer(session);
    }

    // 关闭套接字
    if (session->sock >= 0)
    {
        close(session->sock);
        session->sock = -1;
    }

    // 释放SM2密钥对
    sm2_free_key_pair(&session->local_key);

    // 清除敏感数据
    secure_clear_key(session->shared_key, sizeof(session->shared_key));
    secure_clear_key(session->token, sizeof(session->token));

    // 重置会话状态
    session->state = SESSION_STATE_DISCONNECTED;
    session->user_id = 0;
    memset(session->username, 0, sizeof(session->username));
    memset(session->remote_public_x, 0, sizeof(session->remote_public_x));
    memset(session->remote_public_y, 0, sizeof(session->remote_public_y));

    // 重置传输状态
    memset(&session->current_transfer, 0, sizeof(session->current_transfer));
    session->current_transfer.direction = TRANSFER_DIRECTION_NONE;
    session->current_transfer.state = TRANSFER_STATE_NONE;

    return 0;
}

// 清理客户端
void client_cleanup(void)
{
    // 清理GmSSL资源
    cleanup_gmssl();
}

// 用户注册
int client_register(client_session_t *session, const char *username, const char *password)
{
    if (!session || !username || !password || session->sock < 0 || session->state != SESSION_STATE_KEY_EXCHANGED)
    {
        fprintf(stderr, "[Client Register] Error: Invalid session state (%d) or parameters. Expected KEY_EXCHANGED (%d).\n", session ? session->state : -1, SESSION_STATE_KEY_EXCHANGED);
        return -1;
    }

    // 构造注册请求
    register_request_t request;
    memset(&request, 0, sizeof(request));
    strncpy(request.username, username, sizeof(request.username) - 1);
    // 默认邮箱可以留空

    /* REMOVE CLIENT-SIDE SALT GENERATION AND HASHING
    // 生成随机盐值
    unsigned char salt[SALT_LENGTH];
    if (generate_auth_salt(salt, SALT_LENGTH) != 0)
    {
        fprintf(stderr, "生成盐值失败\n");
        return -1;
    }
    // Copy the generated salt to the request structure
    memcpy(request.salt, salt, SALT_LENGTH);

    // 计算密码哈希（使用SM3哈希函数）
    int hash_op_result = sm3_hash_password(password, salt, SALT_LENGTH,
                                           (unsigned char *)request.password, sizeof(request.password));

    if (hash_op_result != SM3_DIGEST_LENGTH)
    {
        fprintf(stderr, "计算密码哈希失败 (错误码或长度不匹配: %d)\n", hash_op_result);
        return -1;
    }
    */

    // SEND PLAINTEXT PASSWORD (as per design document where server handles hashing+salting)
    strncpy(request.password, password, sizeof(request.password) - 1);
    request.password[sizeof(request.password) - 1] = '\0';

    // Client-side salt is no longer generated or sent; server should generate its own.
    // If request.salt field must exist due to struct, it will be zeroed by memset or ignored by server.

    // fprintf(stderr, "[CLIENT_REGISTER] Attempting registration for user: '%s' on sock %d\n", username, session->sock);
    // 发送注册请求
    uint16_t seq = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_REGISTER_REQUEST, seq, &request, sizeof(request)) != 0)
    {
        fprintf(stderr, "发送注册请求失败\n");
        client_disconnect(session);
        return -1;
    }

    // 接收注册响应
    uint8_t type;
    uint16_t recv_seq;
    void *data = NULL;
    uint32_t len = 0;

    int recv_ret_reg = recv_specific_message(session->sock, MSG_TYPE_REGISTER_RESPONSE, seq, &data, &len);

    if (recv_ret_reg != 0 || data == NULL)
    {
        fprintf(stderr, "接收注册响应失败\n");
        client_disconnect(session);
        return -1;
    }

    if (len != sizeof(register_response_t))
    {
        fprintf(stderr, "接收到的注册响应数据长度不正确\n");
        free(data);
        client_disconnect(session);
        return -1;
    }

    // 处理注册响应
    register_response_t *response = (register_response_t *)data;
    int result = response->result;

    // 如果注册成功，进行自动登录（需要另外调用认证函数）
    free(data);
    return result;
}

// 删除文件
int client_delete_file(client_session_t *session, const char *remote_path)
{
    if (!session || !remote_path || session->sock < 0 || session->state != SESSION_STATE_AUTHENTICATED)
    {
        fprintf(stderr, "[Client Delete] Error: Invalid session state or parameters.\n");
        return -1;
    }

    // 构造删除文件请求
    file_delete_request_t request;
    memset(&request, 0, sizeof(request));
    strncpy(request.path, remote_path, sizeof(request.path) - 1);
    memcpy(request.token, session->token, sizeof(request.token));

    // 发送删除文件请求
    uint16_t seq_sent = session->current_seq++;
    if (send_message(session->sock, MSG_TYPE_FILE_DELETE_REQUEST, seq_sent, &request, sizeof(request)) != 0)
    {
        fprintf(stderr, "发送删除文件请求失败\n");
        // 不在此处调用 client_disconnect, 让上层API的会话恢复机制处理
        return ERR_CONNECTION_ERROR; // 返回连接错误
    }

    // 接收删除文件响应
    void *data = NULL;
    uint32_t len = 0;
    int recv_ret = recv_specific_message(session->sock, MSG_TYPE_FILE_DELETE_RESPONSE, seq_sent, &data, &len);

    if (recv_ret != 0)
    {
        fprintf(stderr, "接收删除文件响应失败，错误码: %d\n", recv_ret);
        // 如果recv_ret是 ERR_SERVER_OPERATION_FAILED，则data可能包含错误信息，但recv_specific_message应已处理
        // 对于其他错误，如连接错误、协议错误，data通常为NULL或应被忽略
        if (data)
            free(data);  // 确保在出错时释放可能的data
        return recv_ret; // 返回从recv_specific_message获取的错误码
    }

    if (len != sizeof(file_delete_response_t) || !data)
    {
        fprintf(stderr, "接收到的文件删除响应数据长度不正确 (len %u, expected %zu) 或数据为空\n", len, sizeof(file_delete_response_t));
        if (data)
            free(data);
        // 这是一个协议错误，可能需要更强的会话重置
        return ERR_PROTOCOL_VERSION; // 或其他合适的协议错误码
    }

    file_delete_response_t *response = (file_delete_response_t *)data;
    uint8_t result_code = response->result;
    free(data);

    if (result_code == ERR_SUCCESS)
    {
        return 0; // Success
    }
    else
    {
        fprintf(stderr, "服务器报告删除失败，错误码: %d\n", result_code);
        return (int)result_code; // 返回服务器的错误码
    }
}